# src/predict.py
import os
import pandas as pd
import matplotlib

matplotlib.use('Agg')  # 设置matplotlib后端避免显示问题
import matplotlib.pyplot as plt
import datetime
from utils.log import Logger
from lightgbm import LGBMClassifier
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
import joblib
# 导入数据库模块
from utils.database import DatabaseConnection
# 导入Redis客户端
from utils.redis_client import RedisClient
# 导入配置管理
from utils.config import ConfigManager

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15
pd.set_option("display.max_columns", None)


class TalentAttritionPredict:
    def __init__(self):
        """
        01- 日志对象
        02- 加载预处理的数据
        03- 加载模型
        """
        # 创建配置管理器
        self.config_manager = ConfigManager()

        # 01- 日志对象
        predict_log_name = "predict_" + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        self.logger = Logger("../", predict_log_name).get_logger()

        # 获取数据库配置
        self.db_config = self.config_manager.get_database_config()

        # 获取模型配置
        self.model_config = self.config_manager.get_model_config()

        # 创建Redis客户端
        redis_config = self.config_manager.get_redis_config()
        self.redis_client = RedisClient(redis_config)

        # 创建数据库连接对象
        self.db_connection = DatabaseConnection(self.db_config, self.redis_client)

        # 02- 加载预处理的数据(加载测试集数据)
        self.data_source = self.load_data_from_db()
        # 03- 加载模型
        model_path = self.model_config.get('model_path', '../model/lightgbm_model.pkl')
        self.logger.info(f'成功加载模型: {model_path}')
        self.model = joblib.load(model_path)

    def load_data_from_db(self):
        """
        从数据库加载测试数据
        """
        try:
            query = "SELECT * FROM test_data"
            # 使用Redis缓存，缓存键为test_data，过期时间为1小时
            data = self.db_connection.execute_query(query, cache_key="test_data", cache_expire=3600)

            self.logger.info(f'从数据库成功加载 {len(data)} 条测试数据')
            return data
        except Exception as e:
            self.logger.error(f'从数据库加载数据失败: {str(e)}')
            raise

    def feature_engineering(self, data):
        """
        特征工程
        """
        data = data.copy()

        # 删除无用特征（先检查列是否存在）
        columns_to_drop = ['id', 'created_at']
        # 检查employee_number列是否存在（考虑大小写不同的情况）
        if 'employee_number' in data.columns:
            columns_to_drop.append('employee_number')
        elif 'EmployeeNumber' in data.columns:
            columns_to_drop.append('EmployeeNumber')

        # 只删除存在的列
        existing_columns = [col for col in columns_to_drop if col in data.columns]
        if existing_columns:
            data = data.drop(existing_columns, axis=1)
            self.logger.info(f"删除了列: {existing_columns}")

        # 处理attrition列（如果存在）
        attrition_column = None
        for col in data.columns:
            if col.lower() == 'attrition':
                attrition_column = col
                break

        if attrition_column:
            data = data.drop(attrition_column, axis=1)

        # 对类别特征进行独热编码
        categorical_features = ['business_travel', 'department', 'education_field', 'gender', 'job_role',
                                'marital_status', 'over_time']
        # 处理类别特征列名大小写问题
        actual_categorical_features = []
        for feature in categorical_features:
            for col in data.columns:
                if col.lower() == feature.lower():
                    actual_categorical_features.append(col)
                    break

        if actual_categorical_features:
            data = pd.get_dummies(data, columns=actual_categorical_features)

        return data

    def evaluate_plot(self, evaluate_df):
        """
        绘制真实标签和预测标签对比的柱状图
        :return:
        """
        plt.figure(figsize=(10, 6))
        real_counts = evaluate_df["真实标签"].value_counts()
        pred_counts = evaluate_df["预测标签"].value_counts()

        x = range(len(real_counts.index))
        width = 0.35

        plt.bar([i - width / 2 for i in x], real_counts.values, width, label='真实标签')
        plt.bar([i + width / 2 for i in x], pred_counts.values, width, label='预测标签')

        plt.xlabel("标签")
        plt.ylabel("数量")
        plt.title("真实标签和预测标签对比")
        plt.xticks(x, real_counts.index)
        plt.legend()
        # 保存图片而不是直接显示
        plt.savefig("人才流失预测结果.png")
        plt.close()  # 关闭图形避免内存泄漏
        self.logger.info("预测结果图已保存为 人才流失预测结果.png")

    def model_predict(self):
        """
        离线预测
        :return: evaluate_df
        """
        self.logger.info("=============开始进行离线预测================")

        # 特征工程
        data = self.data_source.copy()
        X = self.feature_engineering(data)
        y_true = self.data_source['attrition']

        # 模型预测
        y_pred = self.model.predict(X)
        y_pred_proba = self.model.predict_proba(X)[:, 1]

        evaluate_df = pd.DataFrame({
            "真实标签": y_true,
            "预测标签": y_pred,
            "预测概率": y_pred_proba
        })

        # 评估模型
        accuracy = accuracy_score(y_true, y_pred)
        f1 = f1_score(y_true, y_pred)
        auc = roc_auc_score(y_true, y_pred_proba)

        # 绘制真实结果和预测结果对比图像
        self.evaluate_plot(evaluate_df)
        self.logger.info(
            f"=============离线预测结束: Accuracy={accuracy:.4f}, F1 Score={f1:.4f}, AUC={auc:.4f}================")
        self.logger.info("预测结果图已保存为 人才流失预测结果.png")

        # 将评估结果存入Redis
        evaluation_result = {
            "accuracy": accuracy,
            "f1_score": f1,
            "auc": auc,
            "timestamp": datetime.datetime.now().isoformat(),
            "record_count": len(evaluate_df)
        }
        self.redis_client.set_json("latest_evaluation", evaluation_result, ex=86400)  # 缓存24小时

        return evaluate_df


if __name__ == '__main__':
    tap = TalentAttritionPredict()
    tap.model_predict()
